import json
import os
from time import time

import joblib
import pandas as pd
import numpy as np
import stomp
from sklearn.ensemble import GradientBoostingRegressor


class MachineTaskListener(object):
    def __init__(self, host, port):
        self.conn = stomp.Connection10([(host, port)], auto_content_length=False)
        self.conn.connect()
        model_file_path = r'D:\WorkSpace\MQResearch\MQPy\ActiveMqDemo\RpcDemo'
        scaler_file = os.path.join(model_file_path, 'Scaler.model')
        self.scaler = joblib.load(scaler_file)
        model_file = os.path.join(model_file_path, 'Model.model')
        model = joblib.load(model_file)
        self.model = joblib.load(model_file)

    def on_message(self, headers, msg):
        t0 = time()
        print("received message %s" % json.loads(msg))
        task_message = json.loads(msg)
        task_param = json.loads(task_message['TaskParam'])
        x_test = pd.DataFrame(np.array(task_param['Sample']).reshape(-1, 54))

        scaled_df = self.scaler.transform(x_test)
        result = self.model.predict(scaled_df)
        print(x_test)
        print(result)
        task_message['TaskResult'] = result[0]
        task_message['TotalSeconds'] = time() - t0

        replay_msg = json.dumps(task_message)
        self.conn.send(headers['reply-to'], replay_msg)

    def on_error(self, headers, message):
        print('received an error %s' % message)
